import os
import pandas as pd

from preprocessing.option import parse_args


def generate_csv(args, mode):
    """
    :param args:
    :param mode: train, dev, test
    :return:
    """
    if mode == 'train_real' or mode == 'train_fake':
        video_names = get_video_names(args, 'train')
        videos_path = os.path.join(args.root_dir, '{}_files_crops'.format("train".capitalize()))
    else:
        video_names = get_video_names(args, mode)
        videos_path = os.path.join(args.root_dir, '{}_files_crops'.format(mode.capitalize()))
    fold_data = []
    for video_name in video_names:
        video_path = os.path.join(videos_path, video_name)
        for image in os.listdir(video_path):
            video_name_split = video_name.split('_')
            assert len(video_name_split) == 4
            if 'real' in mode or 'fake' in mode:
                if 'real' in mode and int(video_name_split[-1]) == 1:
                    label = 0
                    fold_data.append([video_name, image, label])
                elif 'fake' in mode and int(video_name_split[-1]) != 1:
                    label = 1
                    fold_data.append([video_name, image, label])
            else:
                label = 0 if int(video_name_split[-1]) == 1 else 1
                fold_data.append([video_name, image, label])
    columns = ["video", "file", "label"]
    save_path = '../data/{}/data_{}_{}.csv'.format(args.protocol, args.protocol, mode)
    # save_path = '../data/{}/data_{}_{}_1.csv'.format(args.protocol, args.protocol, mode)

    pd.DataFrame(fold_data, columns=columns).to_csv(save_path, index=False)


def get_video_names(args, mode='dev'):
    video_names = []
    lines = open(os.path.join(args.root_dir, 'Protocols', args.protocol, '{}.txt'.format(mode.capitalize())))
    # lines = open(os.path.join(args.root_dir, 'Protocols', args.protocol, '{}_1.txt'.format(mode.capitalize())))

    for line in lines:
        line_split = line.split(',')
        assert len(line_split) == 2
        video_names.append(line_split[1].strip())
    return video_names


def main():
    args = parse_args()
    protocol_path = os.path.join('../data/{}'.format(args.protocol))
    os.makedirs(protocol_path, exist_ok=True)

    # generate_csv(args, 'train')
    # generate_csv(args, 'dev')
    # generate_csv(args, 'test')

    generate_csv(args, 'train_real')
    generate_csv(args, 'train_fake')


if __name__ == '__main__':
    main()
